import json
import secrets
import rsa
import os
import sys
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../')))
import backend.DID.Key_Management as Key_Management 
import backend.DID.BlockChain as BlockChain 
import backend.DID.DataBase as DataBase 
import backend.DID.En_Decrypt_Management as En_Decrypt_Management
import socket



class DIDSocket:
    def __init__(self, own_did):
        self.own_did = own_did
        self.private_key = DataBase.private_key_changes(own_did)  # 本地私钥
        self.server_socket = None  # 服务器监听socket
        self.connections = []  # 存储所有活跃连接

    def search_publickey(self, did):
        """根据DID从区块链查询公钥"""
        id_ = {'id': f'{did}'}
        table = 'blockchain'
        search_data = 'publickey'

        publickey_bytes = BlockChain.query_data(table, id_, search_data).encode('utf-8')
        return rsa.PublicKey.load_pkcs1(publickey_bytes)

    def recvall(self, sock, length):
        """按指定字节数完整接收数据"""
        data = b''
        while len(data) < length:
            part = sock.recv(length - len(data))
            if not part:
                raise ConnectionError("socket closed")
            data += part
        return data

    def send_with_len(self, sock, data):
        """前置4字节长度后发送数据"""
        sock.sendall(len(data).to_bytes(4, 'big') + data)

    def recv_with_len(self, sock):
        """先读4字节长度，再读正文"""
        header = self.recvall(sock, 4)
        body_len = int.from_bytes(header, 'big')
        return self.recvall(sock, body_len)

    def listen(self, port, host='0.0.0.0'):
        """创建服务器监听socket"""
        self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.server_socket.bind((host, port))
        self.server_socket.listen(5)  # 允许最多5个等待连接
        print(f"DID {self.own_did} 在端口 {port} 监听连接...")
        return self.server_socket

    def accept(self):
        """接受新连接并完成密钥协商"""
        if not self.server_socket:
            raise RuntimeError("尚未调用listen()方法创建监听")

        conn, addr = self.server_socket.accept()
        print(f"收到来自 {addr} 的连接请求")

        # 进行密钥协商
        connection = self._complete_server_handshake(conn)
        if connection:
            self.connections.append(connection)
            return connection, addr
        return None, addr

    def _complete_server_handshake(self, conn):
        """作为服务器完成密钥协商"""
        try:
            # 互换DID
            receiver_did = self.recv_with_len(conn).decode()
            print(f"接收方DID：{receiver_did}")
            self.send_with_len(conn, self.own_did.encode())

            # 接收密文（共享密钥 + random_a）
            receiver_public_key = self.search_publickey(receiver_did)
            encrypt_data = self.recv_with_len(conn)
            decrypt_data = En_Decrypt_Management.decrypt_by_privatekey(encrypt_data, self.private_key)
            if isinstance(decrypt_data, str):
                decrypt_data = decrypt_data.encode('utf-8')

            key_len = 16 if len(decrypt_data) == 32 else len(decrypt_data) - 16
            shared_key = decrypt_data[:key_len]
            random_a = decrypt_data[key_len:key_len + 16]

            # 生成random_b并回包
            random_b = secrets.token_bytes(16)
            data_to_send = random_a + random_b
            encrypt_response = En_Decrypt_Management.encrypt_by_sharekey(shared_key, data_to_send)
            self.send_with_len(conn, encrypt_response)

            # 验证random_b
            encrypted_b = self.recv_with_len(conn)
            decrypted_b = En_Decrypt_Management.decrypt_by_sharekey(shared_key, encrypted_b)
            if isinstance(decrypted_b, str):
                decrypted_b = decrypted_b.encode('utf-8')

            if decrypted_b == random_b:
                print("密钥协商成功")
                print(f"共享密钥: {shared_key.hex()}")
                return {
                    'socket': conn,
                    'peer_did': receiver_did,
                    'peer_public_key': receiver_public_key,
                    'shared_key': shared_key,
                    'address': conn.getpeername()
                }
            else:
                print("密钥协商失败，随机数不同")
                conn.close()
                return None

        except Exception as e:
            print(f"DID密钥协商失败：{e}")
            conn.close()
            return None

    def connect(self, receiver_did, port, host):
        """连接到远程DID并完成密钥协商"""
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.connect((host, port))
        print(f"连接到 {host}:{port}，对方DID:{receiver_did}")

        # 完成客户端密钥协商
        connection = self._complete_client_handshake(sock, receiver_did)
        if connection:
            self.connections.append(connection)
            return connection
        return None

    def _complete_client_handshake(self, sock, receiver_did):
        """作为客户端完成密钥协商"""
        try:
            # 互换DID
            self.send_with_len(sock, self.own_did.encode())
            received_did = self.recv_with_len(sock).decode()
            print(f"收到对方DID: {received_did}")
            if receiver_did != received_did:
                print(f"DID不匹配: 期望{receiver_did}, 实际{received_did}")
                sock.close()
                return None

            # 发送（共享密钥 + random_a）的密文
            receiver_public_key = self.search_publickey(receiver_did)
            shared_key = Key_Management.generate_aes_key()
            random_a = secrets.token_bytes(16)
            data_to_encrypt = shared_key + random_a
            encrypted_data = En_Decrypt_Management.encrypt_by_publickey(data_to_encrypt, receiver_public_key)
            self.send_with_len(sock, encrypted_data)

            # 接收random_a + random_b
            encrypted_response = self.recv_with_len(sock)
            decrypted_response = En_Decrypt_Management.decrypt_by_sharekey(shared_key, encrypted_response)
            if isinstance(decrypted_response, str):
                decrypted_response = decrypted_response.encode('utf-8')

            key_len = len(shared_key)
            received_a = decrypted_response[:16]  # random_a固定16字节
            random_b = decrypted_response[16:32]

            if received_a != random_a:
                print("密钥协商失败: 随机数A验证不通过")
                sock.close()
                return None

            # 回传random_b
            encrypted_b = En_Decrypt_Management.encrypt_by_sharekey(shared_key, random_b)
            self.send_with_len(sock, encrypted_b)
            print("密钥协商成功")
            print(f"共享密钥: {shared_key.hex()}")

            return {
                'socket': sock,
                'peer_did': receiver_did,
                'peer_public_key': receiver_public_key,
                'shared_key': shared_key,
                'address': sock.getpeername()
            }

        except Exception as e:
            print(f"DID密钥协商失败：{e}")
            sock.close()
            return None

    def send(self, connection, meg):
        """通过指定连接发送加密数据"""
        try:
            data=json.dumps(meg).encode('utf-8')
            sock = connection['socket']
            shared_key = connection['shared_key']

            if not sock or not shared_key:
                print("发送失败: 无效的连接或共享密钥")
                return False

            # 确保数据是字节类型
            if isinstance(data, str):
                data = data.encode('utf-8')

            encrypted_data = En_Decrypt_Management.encrypt_by_sharekey(shared_key, data)
            self.send_with_len(sock, encrypted_data)
            print(f"向 {connection['peer_did']} 发送数据成功")
            return True

        except Exception as e:
            print(f"发送数据失败: {e}")
            self.close_connection(connection)
            return False

    def recv(self, connection):
        """通过指定连接接收并解密数据"""
        try:
            sock = connection['socket']
            shared_key = connection['shared_key']

            if not sock or not shared_key:
                print("接收失败: 无效的连接或共享密钥")
                return None

            encrypted_data = self.recv_with_len(sock)
            if not encrypted_data:
                print("未接收到数据")
                return None

            return json.loads(En_Decrypt_Management.decrypt_by_sharekey(shared_key, encrypted_data).decode())

        except Exception as e:
            print(f"接收数据失败: {e}")
            self.close_connection(connection)
            return None

    def close_connection(self, connection):
        """关闭指定连接"""
        try:
            if connection in self.connections:
                self.connections.remove(connection)

            sock = connection.get('socket')
            if sock:
                sock.close()

            print(f"与 {connection.get('peer_did')} 的连接已关闭")

        except Exception as e:
            print(f"关闭连接失败: {e}")

    def close(self):
        """关闭所有连接和服务器socket"""
        # 关闭所有活跃连接
        for connection in self.connections[:]:
            self.close_connection(connection)

        # 关闭服务器监听socket
        if self.server_socket:
            self.server_socket.close()
            self.server_socket = None
            print("服务器监听已关闭")

        print("所有连接已关闭，资源已释放")
